/*
 *https://leetcode.cn/problems/count-pairs-of-connectable-servers-in-a-weighted-tree-network/description/
 *3067. 在带权树网络中统计可连接服务器对数目
 *medium 吴朝泽 2024.10.26
 *DFS深度优先搜索
*/

class Solution {
public:
    vector<int> countPairsOfConnectableServers(vector<vector<int>>& edges, int signalSpeed) {
        int n = edges.size() + 1; // 服务器数量为边的数量加一
        vector<pair<int, int>> g[n]; // 创建邻接表

        // 构建图
        for (auto& e : edges) {
            int a = e[0], b = e[1], w = e[2]; // 获取边的两个端点和权重
            g[a].emplace_back(b, w); // 在图中添加边
            g[b].emplace_back(a, w); // 因为是无向图，反向边也要添加
        }

        // 深度优先搜索（DFS）函数
        function<int(int, int, int)> dfs = [&](int a, int fa, int ws) {
            int cnt = ws % signalSpeed == 0; // 如果当前路径的信号速度能连接，则计数加一
            for (auto& [b, w] : g[a]) { // 遍历与当前节点相连的所有边
                if (b != fa) { // 确保不回到父节点
                    cnt += dfs(b, a, ws + w); // 递归调用DFS，累加路径权重
                }
            }
            return cnt; // 返回可连接的服务器数量
        };

        vector<int> ans(n); // 存储每个服务器的可连接对数量
        for (int a = 0; a < n; ++a) { // 遍历每个服务器
            int s = 0; // 当前可连接的服务器数量
            for (auto& [b, w] : g[a]) { // 遍历与当前服务器相连的所有边
                int t = dfs(b, a, w); // 从当前服务器出发进行DFS
                ans[a] += s * t; // 更新当前服务器的可连接对数量
                s += t; // 更新可连接服务器的总数
            }
        }
        return ans; 
    }
};
